from fastapi import FastAPI, File, UploadFile
import shutil
import os
import cv2
import numpy as np
import pandas as pd
import logging
import uvicorn
from skimage.measure import label, regionprops
from fastapi.middleware.cors import CORSMiddleware
from time import time
from typing import List
from fastapi.staticfiles import StaticFiles
app = FastAPI()
app.mount("/static", StaticFiles(directory="OCT_layer_cal/image/processed_results"), name="static")
app.mount("/colorinput", StaticFiles(directory="user_color_data"), name="colorinput")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源跨域
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ImagePath:
    def __init__(self, path: str):
        self.path = path

logging.basicConfig(level=logging.DEBUG)


def extract_center_line(binary_image):
    # 使用传统的细化算法进行中心线提取
    element = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    done = False
    size = np.size(binary_image)
    skel = np.zeros(binary_image.shape, np.uint8)

    while not done:
        eroded = cv2.erode(binary_image, element)
        temp = cv2.dilate(eroded, element)
        temp = cv2.subtract(binary_image, temp)
        skel = cv2.bitwise_or(skel, temp)
        binary_image = eroded.copy()

        zeros = size - cv2.countNonZero(binary_image)
        if zeros == size:
            done = True

    return skel

def calculate_average_width(center_line_points, contours):
    average_widths = []

    for point in center_line_points:
        distances = []
        for contour in contours:
            distance = cv2.pointPolygonTest(contour, (int(point[1]), int(point[0])), True)
            distances.append(abs(distance))

        # 计算中心线到轮廓的最小距离
        min_distance = min(distances)

        # 将最小距离乘以2，即为连通域在该点的宽度
        width_at_point = min_distance * 2
        average_widths.append(width_at_point)

    # 计算平均宽度
    average_width = np.mean(average_widths)
    return average_width


def process_image(image_path, output_path_csv, color, slice):
    # 读取图像并进行二值化处理
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    _, binary_image = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

    # 提取中心线
    center_line = extract_center_line(binary_image)
    # _, contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    # 本地必须用下面这个
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)


    # 计算平均宽度
    center_line_points = np.column_stack(np.where(center_line > 0))
    average_width = calculate_average_width(center_line_points, contours)

    # 初始化layer变量为默认值
    layer = 'Unknown Layer'

    # 根据颜色值设置layer
    if all(color == [0, 0, 255]):
        layer = 'ILM-IPL'
    elif all(color == [0, 255, 0]):
        layer = 'IPL-OPL'
    elif all(color == [0, 255, 192]):
        layer = 'OPL-ISOS'
    elif all(color == [156, 255, 0]):
        layer = 'ISOS-RPE'
    elif all(color == [255, 0, 0]):
        layer = 'RPE-BM'

    data = {'Slice': [slice], 'Layer': [layer], 'Average Height (px)': [average_width]}
    df = pd.DataFrame(data)

    # 获取输出文件夹名和csv文件名
    folder_name = output_path_csv.split('/')[-1]
    csv_file_name = f"{folder_name}.csv"

    # 构建完整的输出CSV路径
    full_output_path = f"{output_path_csv}/{csv_file_name}"
    # 保证文件夹一定存在
    os.makedirs(os.path.dirname(full_output_path), exist_ok=True)
    df.to_csv(full_output_path, mode='a', index=False, header=(not os.path.isfile(full_output_path)))


@app.post("/upload_image/")
async def upload_image(image: UploadFile = File(...)):
    timestamp = int(time())
    os.makedirs(f"{os.getcwd()}/user_color_data/{timestamp}", exist_ok=True)
    save_path = f"{os.getcwd()}/user_color_data/{timestamp}/{image.filename}"
    with open(save_path, "wb") as img:
        shutil.copyfileobj(image.file, img)

    return {"message": "Image uploaded successfully","time":timestamp}



@app.post("/process_image/")
async def process_uploaded_image(user_string: int):
    input_path = f'{os.getcwd()}/user_color_data/{user_string}'
    # print(input_path)
    output_path = f'{os.getcwd()}/OCT_layer_cal'
    # print(output_path)
    os.makedirs(output_path, exist_ok=True)

    # 遍历原始图像文件夹
    for image_file in os.listdir(input_path):
        if image_file.endswith('.png') or image_file.endswith('.jpg'):
            # 构建完整的输入图像路径
            image_path = os.path.join(input_path, image_file)

            # 定义结果文件夹名称和切片图像名称
            folder_name = 'processed_results'
            slice_image_name = f'{user_string}'
            # 处理当前图像
            color_label_take(image_path, output_path, folder_name, slice_image_name)
            print(f"Image {slice_image_name} processed successfully")

    return {"processed_image": f"Image {os.listdir(input_path)} processed successfully"}


@app.get("/list_images/")
async def list_images(user_string: int):
    image_paths = []
    folder_path = f'{os.getcwd()}/OCT_layer_cal/image/processed_results/{user_string}/'
    # 遍历指定文件夹中的所有文件
    for file_name in os.listdir(folder_path):
        file_path = os.path.join(f'http://127.0.0.1:8000/static/{user_string}/', file_name)

        # 检查文件是否为图片文件（这里仅检查扩展名）
        if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
            image_paths.append(file_path)

    return image_paths

def color_label_take(label, output_path, folder_name, slice_image_name):
    label = cv2.imread(label)
    # 获取图像形状
    height, width, _ = label.shape

    # 提取颜色信息
    unique_colors = np.unique(label.reshape(-1, label.shape[2]), axis=0)
    unique_colors = unique_colors[unique_colors.any(axis=1)]
    # 确保输出路径存在
    output_path_csv = os.path.join(output_path, 'csv', folder_name)
    output_path_image = os.path.join(output_path, 'image', folder_name, slice_image_name)
    os.makedirs(output_path_image, exist_ok=True)
    os.makedirs(output_path_csv, exist_ok=True)

    for color in unique_colors:
        # 创建一个与原始图像相同形状的零数组
        result_array_label = np.zeros((height, width, label.shape[2]), dtype=np.uint8)

        # 找出当前颜色的像素位置
        color_mask = (label == color).all(axis=-1)

        # 对于3通道图像，进行连通域分析（需要先转为灰度图）
        gray_color_mask = color_mask.sum(axis=-1) > 0
        _, cc_labels, stats, _ = cv2.connectedComponentsWithStats(gray_color_mask.astype(np.uint8), connectivity=4)

        # 找到面积最大的连通域
        max_area = -1
        max_label = None
        for i, prop in enumerate(regionprops(cc_labels)):
            if prop.area > max_area and prop.area > 0:
                max_area = prop.area
                max_label = i + 1  # 连通域标签从1开始

        # 将最大连通域内的像素恢复成当前颜色
        if max_label is not None:
            result_array_label[color_mask & (cc_labels == max_label)] = color

        # 保存每个颜色的最大连通域结果
        cv2.imwrite(f'{output_path_image}/{str(slice_image_name)}_{tuple(color)}.jpg', result_array_label)

        process_image(f'{output_path_image}/{str(slice_image_name)}_{tuple(color)}.jpg', output_path_csv, color,
                      slice_image_name)
        # print(f'oct厚层分析完成---{slice_image_name}')


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000)
